#include <cstdlib>
#include <cstring>
#include <iostream>

#include <boost/array.hpp>
#include <boost/foreach.hpp>

#include <alloca.h>
#include <stdio.h>
#include <errno.h>

#include "c_connection.h"


namespace dhd
{
namespace client
{

   using boost::asio::mutable_buffer;

   using boost::asio::detail::socket_ops::network_to_host_short;


   Connection::Connection( ServerCFG * cfg, IConnectionHandler *_handler, boost::asio::io_service *io_svc ) :
   conf(cfg), io_service(io_svc), ssl_ctx(0), sock(0), ssl_sock(0), sock_io(0), srp_usr(0), handler(_handler),
   reconnect_timer(*io_svc), connected(false), got_server_root(false), authenticated(false), srp_sent_m(false)
   {
       tcp::resolver resolver(*io_service);
        
       char buff[10];
       sprintf(buff, "%d", conf->port);
        
       tcp::resolver::query query(tcp::v4(), conf->hostname.c_str(), buff);
       tcp::resolver::iterator iterator = resolver.resolve(query);
       
       svr_endpoint = *iterator;
    
       if (conf->use_ssl)
       {
           ssl_ctx = new boost::asio::ssl::context(*io_service, boost::asio::ssl::context::tlsv1_client);
           ssl_ctx->set_verify_mode( boost::asio::ssl::context::verify_none );
           SSL_CTX_set_cipher_list( ssl_ctx->impl(), "ADH-AES256-SHA" );
       }
    
       async_connect();
   }
   
   
   Connection::~Connection()
   {
       boost::system::error_code  e;
       
       reconnect_timer.cancel( e );
       
       delete sock;      sock      = 0;
       delete ssl_sock;  ssl_sock  = 0;
       delete ssl_ctx;   ssl_ctx   = 0;
       
       sock_io = 0;
       
       if (srp_usr)
           srp_user_delete(srp_usr);       
   }
   
   
   void Connection::on_disconnect( bool reconnect )
   {
       if (srp_usr)
           srp_user_delete(srp_usr);

       sock->close();
       
       delete sock;
       delete ssl_sock;
       
       sock      = 0;
       ssl_sock  = 0;
       sock_io   = 0;
       srp_usr   = 0;
       
       connected       = false;
       got_server_root = false;
       authenticated   = false;
       srp_sent_m      = false;
       
       if (reconnect)
           async_connect();

       std::cerr << "############ DISCONNECTED #############" << std::endl;
       
       handler->connection_on_disconnect();
   }


   void Connection::async_connect()
   {
       if (!sock)
       {
           sock = new tcp::socket(*io_service);
        
           if (conf->use_ssl)
               ssl_sock = new ssl_socket(*sock, *ssl_ctx);
       }
       
       
       if (conf->use_ssl)
       {
           sock->async_connect(svr_endpoint, boost::bind( &Connection::do_ssl_handshake,
                                                          this,
                                                          boost::asio::placeholders::error) );
       }
       else
       {
           sock->async_connect(svr_endpoint, boost::bind( &Connection::on_async_connect_complete,
                                                          this,
                                                          boost::asio::placeholders::error) );
       }
   }
   

   void Connection::async_reconnect( const boost::system::error_code& error )
   {
       if (error)
           return;  // Reconnect canceled. Exit.
       std::cerr << "Attempting reconnect to " << svr_endpoint << std::endl;
       async_connect();
   }


   void Connection::do_ssl_handshake( const boost::system::error_code& error )
   {
       if (error)
       {
           on_async_connect_complete( error );
       }
       else
       {
           ssl_sock->async_handshake(boost::asio::ssl::stream_base::client,
                                     boost::bind( &Connection::on_async_connect_complete,
                                                  this, 
                                                  boost::asio::placeholders::error) );
       }
   }
   
   
   void Connection::on_async_connect_complete( const boost::system::error_code& error )
   {
       if (error)
       {
           reconnect_timer.expires_from_now(boost::posix_time::seconds(5));
           reconnect_timer.async_wait( boost::bind( &Connection::async_reconnect, 
                                                    this, 
                                                    boost::asio::placeholders::error) );
           return;
       }
       
       sock_io = new SocketIO( sock, ssl_sock, 
                               boost::bind( &Connection::handle_message,  this, _1, _2, _3 ),
                               boost::bind( &Connection::on_read_error,   this, _1, _2 ) );
       
       if ( conf->use_srp )
       {
           dhd::ToServer to_svr;
           
           const char *          ignore;
           const unsigned char * bytes_A;
           int                   len_A;
           
           srp_usr = srp_user_new( conf->frontend_cfg->srp_username.c_str(), 
                                   (const unsigned char *) conf->frontend_cfg->srp_password.data(), 
                                   conf->frontend_cfg->srp_password.length() );
           
           srp_user_start_authentication( srp_usr, &ignore, &bytes_A, &len_A );
           
           std::string A( (const char *) bytes_A, len_A );
           
           to_svr.set_msg_type( dhd::DHD_AUTHENTICATION );
           
           ::dhd::Authentication* a = to_svr.mutable_auth();
           
           a->set_username( conf->frontend_cfg->srp_username );
           a->set_a(A);
           
           PBuffer pbuff = allocate_buffer( to_svr.ByteSize() );
           
           to_svr.SerializeToArray( pbuff->data(), pbuff->size() );
           
           sock_io->enqueue( 0, pbuff, 0 );
       }
       else
           authentication_complete();
   }
   

   void Connection::authentication_complete()
   {
       authenticated = true;
       connected     = true;
       
       std::cerr << "############ CONNECTED #############" << std::endl;
 
       handler->connection_on_connect();
   }
   
   
   void Connection::on_read_error( const char * msg, bool reconnect )
   {
       std::cerr << "Error: " << msg << std::endl;
       on_disconnect( reconnect );
   }
   
   
   void Connection::handle_message( uint16_t request_id, PBuffer header, PBuffer buffer )
   {
       std::cerr << "######## GOT MSG FROM SERVER with req_id:  " << request_id << std::endl;
       
       if ( request_id == 0 )
       {
           ToClient tc;
           
           if ( !tc.ParseFromArray( header->data(), header->size() ))
               throw BadProtocol("Failed to parse reply message");

           if ( authenticated )
           {
               handler->connection_on_unsolicited( tc );
           }
           else if ( tc.msg_type() == dhd::DHD_AUTHENTICATION )
           {
               handle_authentication_message( tc );
           }
           else
               throw BadProtocol("Invalid message received during authentication");               
       }
       else
       {
           ReplyMap::iterator itor = pending_replies.find( request_id );
           
           if ( itor != pending_replies.end() ) 
           {
               std::cerr << "##### FOUND PENDING REQUEST ####" << std::endl;
               
               PInteraction pi = itor->second;
               pending_replies.erase( itor );
               
               if ( !pi->reply.ParseFromArray( header->data(), header->size() ))
                   throw BadProtocol("Failed to parse reply message");
               
               assert(authenticated);
               
               pi->pbuffer = buffer;

               handler->connection_on_reply( pi );    
           }
           else
           {
               throw BadProtocol("***** Invalid request_id! Ignoring *****");
           }
       }
   }

   
   void Connection::handle_authentication_message( const ToClient & reply )
   {   
       if ( authenticated )
           throw BadProtocol("Received authentication message on already authenticated connection", false);
       
       if ( !srp_usr )
           throw BadProtocol("Received unexpected and invalid authentication message", false);
       
       if ( !reply.has_auth() )
           throw BadProtocol("Auth packet missing Auth message!", false);
               
       if ( !srp_sent_m )
       {
           // step 2
           if ( !reply.auth().has_s() || !reply.auth().has_b() )
               throw BadProtocol("Auth packet missing s & B!", false);
           
           const unsigned char  *bytes_m;

           srp_user_process_challenge( srp_usr, 
                                       (const unsigned char *)reply.auth().s().data(), reply.auth().s().length(), 
                                       (const unsigned char *)reply.auth().b().data(), reply.auth().b().length(), 
                                       &bytes_m );
               
           if ( !bytes_m )
               throw BadProtocol("Authentication safety check failed", false);

           ToServer               to_svr;
           ::dhd::Authentication *a = to_svr.mutable_auth();
           
           to_svr.set_msg_type( dhd::DHD_AUTHENTICATION );
           
           std::string m( (const char *) bytes_m, SRP_SHA_DIGEST_LENGTH );
           a->set_m( m );
               
           PBuffer pbuff = allocate_buffer( to_svr.ByteSize() );
           
           to_svr.SerializeToArray( pbuff->data(), pbuff->size() );
               
           sock_io->enqueue( 0, pbuff, 0 );
           
           srp_sent_m = true;
       }
       else
       {
           // step 3
           if ( !reply.auth().has_hamk() )
               throw BadProtocol("Authentication packet missing HAMK", false);

           srp_user_verify_session( srp_usr, (const unsigned char *)reply.auth().hamk().data() );
           
           if ( !srp_user_is_authenticated(srp_usr))
               throw BadProtocol("Authentication Failed", false);
           
           std::cerr << "SRP AUTH SUCCESS!!!!" << std::endl;
                   
           authenticated = true;
                                      
           authentication_complete();
       }
   }
   
   
   void Connection::send_interaction( PInteraction i )
   {
       if ( i->reply_required() )
           pending_replies[ i->request_id ] = i;
       
       PBuffer pbuff = allocate_buffer( i->request.ByteSize() );
       
       i->request.SerializeToArray( pbuff->data(), pbuff->size() );
       
       sock_io->enqueue( i->request_id, pbuff, i->pbuffer );
   }

   
} // end namespace dhd::client
} // end namespace dhd

